import cv2
import numpy as np
import os


def main():
    root = r'predict_online'
    with open(os.path.join(root, "label_list.txt"), 'r') as file:
        label_list = file.readlines()
    label_list = [label.rstrip() for label in label_list]  # 去掉空字符
    label_list = [label for label in label_list if label != '']  # 去掉空行

    drawing = False
    x1, y1 = 0, 0
    img = np.zeros((600, 800, 3), np.uint8)

    mouse_param = {}
    mouse_param['ori_pos'] = None
    mouse_param['drawing'] = drawing
    mouse_param['x1'] = x1
    mouse_param['y1'] = y1
    mouse_param['img'] = img
    mouse_param['x2'] = 0
    mouse_param['y2'] = 0
    mouse_param['count'] = 0
    mouse_param['label_list'] = label_list

    cv2.namedWindow('img')
    cv2.setMouseCallback('img', on_mouse, mouse_param)
    while 1:
        cv2.imshow('img', img)
        k = cv2.waitKey(1) & 0xFF
        if k == 27:  # 按esc退出
            break
    cv2.destroyAllWindows()
    # print(os.path.join(root, 'prediction', img_name))
    cv2.imwrite(os.path.join('predict_online', 'result.jpg'), mouse_param['img'], [int(cv2.IMWRITE_JPEG_QUALITY), 100])


def on_mouse(event, x, y, flags, param):
    img = param['img']
    x1 = param['x1']
    y1 = param['y1']
    if event == cv2.EVENT_LBUTTONDOWN:  # 鼠标左键被按下
        param['drawing'] = True
        param['x1'], param['y1'] = x, y
        print('down')
    elif event == cv2.EVENT_MOUSEMOVE and flags == cv2.EVENT_FLAG_LBUTTON:  # 按住鼠标左键拖动
        if param['drawing'] is True:
            pass
            # cv2.rectangle(img, (ix, iy), (x, y), (0, 255, 0), 1)
        # print('move')
    elif event == cv2.EVENT_LBUTTONUP:
        param['drawing'] = False
        param['x2'] = x
        param['y2'] = y
        if param['count'] in range(len(param['label_list'])):
            cv2.rectangle(img, (x1, y1), (x, y), (255, 0, 0), 1)
            cv2.putText(img, param['label_list'][param['count']], (x1, y1), color=(0, 255, 0), fontFace=cv2.FONT_HERSHEY_COMPLEX, fontScale=0.5)
        # 画框计数
        param['count'] += 1
        print('up')


if __name__ == '__main__':
    main()
